{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyP2LmoAi6MCx6sKIv8Ecumw",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/ESMFold_api.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#**ESMFold API**\n",
        "for more details see: [Github](https://github.com/facebookresearch/esm/tree/main/esm), [Preprint](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v1)\n",
        "\n",
        "#### **Tips and Instructions**\n",
        "- click the little ▶ play icon to the left of each cell below.\n",
        "\n",
        "#### **Limitations**\n",
        "- max total length=400\n",
        "- for lengths > 400, try our [colab notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ESMFold.ipynb)\n",
        "---\n",
        "\n",
        "*credit: [@drchristhorpe](https://twitter.com/drchristhorpe)*"
      ],
      "metadata": {
        "id": "v2jU2RLyacXQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "M8p9x4xCXCrw"
      },
      "outputs": [],
      "source": [
        "#@title run **ESMFold**\n",
        "import os, time, re, requests\n",
        "if not os.path.isfile(\"colabfold.py\"):\n",
        "  os.system(\"pip -q install py3Dmol\")\n",
        "  os.system(\"wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/colabfold.py\")\n",
        "  os.system(\"mkdir -p tmp\")\n",
        "\n",
        "import colabfold as cf\n",
        "\n",
        "def predict(sequence):\n",
        "  esmfold_api_url = 'https://api.esmatlas.com/foldSequence/v1/pdb/'\n",
        "  r = requests.post(esmfold_api_url, data=sequence)\n",
        "  while r.status_code != 200:\n",
        "    time.sleep(5)\n",
        "    r = requests.post(esmfold_api_url, data=sequence)\n",
        "  structure = r.text\n",
        "  return structure\n",
        "\n",
        "sequence = \"GWSTELEKHREELKEFLKKEGITNVEIRIDNGRLEVRVEGGTERLKRFLEELRQKLEKKGYTVDIKIE\" #@param {type:\"string\"}\n",
        "sequence = re.sub(\"[^A-Z]\", \"\", sequence.upper())\n",
        "#assert len(sequence) <= 400, \"error: max length supported is 400\"\n",
        "hash = cf.get_hash(sequence)[:5]\n",
        "pdb_filename = f\"tmp/prediction_{hash}.pdb\"\n",
        "if not os.path.isfile(pdb_filename):\n",
        "  pdb_str = predict(sequence)\n",
        "  with open(pdb_filename,\"w\") as out:\n",
        "    out.write(pdb_str)\n",
        "\n",
        "color = \"confidence\" #@param [\"confidence\", \"rainbow\"]\n",
        "if color == \"confidence\": color = \"lDDT\"\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "v = cf.show_pdb(pdb_filename,\n",
        "                show_sidechains,\n",
        "                show_mainchains, color,\n",
        "                color_HP=True,\n",
        "                size=(800,480),\n",
        "                vmin=0.5, vmax=0.9)       \n",
        "v.setHoverable({}, True,\n",
        "               '''function(atom,viewer,event,container){if(!atom.label){atom.label=viewer.addLabel(\"      \"+atom.resn+\":\"+atom.resi,{position:atom,backgroundColor:'mintcream',fontColor:'black'});}}''',\n",
        "               '''function(atom,viewer){if(atom.label){viewer.removeLabel(atom.label);delete atom.label;}}''')\n",
        "v.show()           \n",
        "if color == \"lDDT\":\n",
        "  cf.plot_plddt_legend().show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title download prediction\n",
        "from google.colab import files\n",
        "files.download(pdb_filename)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "fcR0YV4IaRpU"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}